
import tushare as ts
from config import TUSHARE_TOKEN
import pandas as pd

#设置token
ts.set_token(TUSHARE_TOKEN)
pro = ts.pro_api()

def get_stock_daily(ts_code:str,start_date:str , end_date:str ):
    """
    获取单只股票的日线行情数据
    :return: df
    """
    df = pro.daily(ts_code = ts_code , start_date = start_date,end_date = end_date)
    df['trade_date'] = pd.to_datetime(df['trade_date'])
    df.set_index('trade_date',inplace = True)
    df.sort_index(inplace=True)
    return df

def get_index_constituents(ts_code='000300.SH'):
    """
    获取指数成分股（例如沪深300），作为我们的股票池
    """
    df = pro.stock_basic()
    return df['ts_code']


def get_all_data(start_date='20210101', end_date='20241231') -> dict:
    """
    获取股票池内所有股票的日线数据
    """
    # 获取沪深300成分股代码
    stock_list = get_index_constituents()

    all_data = {}
    for ts_code in stock_list[:50]:  # 这里只取50只股票加速回测，实战中可以更多
        print(f"Fetching data for {ts_code}")
        try:
            df = get_stock_daily(ts_code, start_date, end_date)
            # 检查获取的数据是否为空
            if df is not None and not df.empty:
                all_data[ts_code] = df
                print(f"Successfully fetched {len(df)} rows for {ts_code}")
            else:
                print(f"Empty data for {ts_code}")
        except Exception as e:
            print(f"Error fetching {ts_code}: {e}")

    # 最终检查整个数据集是否为空
    if not all_data:
        print("Warning: No data was fetched for any stock!")
    else:
        print(f"Successfully fetched data for {len(all_data)} stocks")

    return all_data





if __name__ == "__main__":

    data = get_all_data()
    # 可以将数据保存到本地，避免每次重新下载
    print(data)
    import pickle
    with open('stock_data.pkl', 'wb') as f:
        pickle.dump(data, f)
